import torch
import torch.nn as nn
from torch.autograd import Variable
import numpy as np
import logging
import os
import pickle

from Config.Config import ACTNET200V13_PKL
from DataSet.Dataset_TSNscore_Pair import Dataset_TSNscore_Pair
from Tools.utils import nms, recall_vs_iou_thresholds, get_gt, convert
from Model.origin_tcn import TCN_Pair2

SAVE_DIR = '/mnt/md1/Experiments/PYTSN_Test2/'
CKPT = '/mnt/md1/Experiments/PYTSN_Test2/pair2_model_17.ckpt'
# SUBSET = 'validation'
SUBSET = 'testing'

model = TCN_Pair2().cuda()
model = nn.DataParallel(model)
model.load_state_dict(torch.load(CKPT))
softmax = nn.Softmax(dim=-1).cuda()

val_dataset = Dataset_TSNscore_Pair(SUBSET, 'test')
gen = val_dataset.enum_genfucntion()

ranked_proposals = dict()
last_vid = None
cnt = 0
for idx,[vid,proposals,feat] in enumerate(gen):

    if len(proposals) == 0:
        continue

    feat = np.expand_dims(feat,1)
    feat = Variable(torch.from_numpy(feat)).cuda().float()
    out = model(feat[:,:,:12,:],feat[:,:,12:,:])
    out = softmax(out)
    prob = out[:,0]

    for i in range(proposals.shape[0]):
        proposals[i,2] = 1 - prob[i]

    if vid not in ranked_proposals.keys():
        ranked_proposals[vid] = []

    ranked_proposals[vid].append(proposals)

    show = False
    if last_vid != vid:
        last_vid = vid
        cnt += 1
        show = True

    if cnt % 100 == 0 and show:
        print('ranked proposal ...',len(ranked_proposals))

with open(os.path.join(SAVE_DIR,'ranked_proposals.pkl'),'wb') as f:
    pickle.dump(ranked_proposals,f)

# with open(os.path.join(SAVE_DIR,'ranked_proposals.pkl'),'rb') as f:
#       ranked_proposals = pickle.load(f)

with open(ACTNET200V13_PKL, 'rb') as f:
    gt = pickle.load(f)['database']

proposal_at_1 = {'s-init':[],'s-end':[],'score':[],'label':[],'video-id':[]}

for vid, proposal in ranked_proposals.items():

    proposal = np.asarray(proposal)
    # 处于未知的原因需要特别判断一下proposal的形状
    if len(proposal.shape) == 3:
        proposal = proposal.reshape((proposal.shape[0]*proposal.shape[1], proposal.shape[2]))
    else:
        proposal = np.concatenate(proposal)

    keep_ind = nms(proposal, proposal[:,2], 0.45)
    proposal = proposal[keep_ind,:]

    proposal_at_1['s-init'].append(proposal[0,0])
    proposal_at_1['s-end'].append(proposal[0,1])
    proposal_at_1['score'].append(proposal[0,2])
    proposal_at_1['video-id'].append('v_'+vid)

print('START EVAL ...')
gt2 = get_gt(gt)
iou_thrs = np.arange(0.1, 1.0, 0.1)
recall_at_1 = recall_vs_iou_thresholds(convert(proposal_at_1), gt2, iou_threshold=iou_thrs)
print(np.array_str(np.vstack([iou_thrs, recall_at_1]), precision=4, suppress_small=True),flush=True)
